import os
import json

import cv2 as cv
from math import *
import math
import random
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt


img_dir1 = 'E:/datasets/client-data/sen12_s100'
img_dir2 = 'E:/datasets/client-data/sen12_s100_aug'
json_file = "E:/workspace/SOMatch/tmp/json/sen12_tt_harris/pt_s100.json"
dst_json_path = "E:/datasets/client-data/sen12_s100_aug/s100_aug.json"

# def rota_rect(box, theta1, x, y):
#     """
#     :param box: 正矩形的四个顶点
#     :param theta: 旋转角度
#     :param x: 旋转中心(x,y)
#     :param y: 旋转中心(x,y)
#     :return: 旋转矩形的四个顶点坐标
#     """
#     # 旋转矩形
#     box_matrix = np.array(box) - np.repeat(np.array([[x, y]]), len(box), 0)
#     theta = theta1 / 180. * np.pi
#     rota_matrix = np.array([[np.cos(theta), -np.sin(theta)],
#                             [np.sin(theta), np.cos(theta)]], np.float)
#
#     new_box = box_matrix.dot(rota_matrix) + np.repeat(np.array([[x, y]]), len(box), 0)
#     return new_box
# # 裁剪最小外接矩形
# def imagecrop(image, box):
#     xs = [x[1] for x in box]
#     ys = [x[0] for x in box]
#     xs_min = int(min(xs) + 0.5)
#     xs_max = int(max(xs) + 0.5)
#     ys_min = int(min(ys) + 0.5)
#     ys_max = int(max(ys) + 0.5)
#     # print(xs)
#     # print(min(xs), max(xs), min(ys), max(ys))
#     cropimage = image[xs_min:xs_max, ys_min:ys_max]
#     # print(cropimage.shape)
#     return cropimage

def rotateImage(img,degree,pt1,pt2,pt3,pt4):
    height, width = img.shape[:2]
    heightNew = int(width * fabs(sin(radians(degree))) + height * fabs(cos(radians(degree))))
    widthNew = int(height * fabs(sin(radians(degree))) + width * fabs(cos(radians(degree))))
    matRotation = cv.getRotationMatrix2D((width/2, height/2), degree, 1)
    matRotation[0, 2] += (widthNew - width) / 2
    matRotation[1, 2] += (heightNew - height) / 2
    imgRotation = cv.warpAffine(img, matRotation, (widthNew, heightNew), borderValue=(255, 255, 255))
    pt1 = list(pt1)
    pt3 = list(pt3)
    [[pt1[0]], [pt1[1]]] = np.dot(matRotation, np.array([[pt1[0]], [pt1[1]], [1]]))
    [[pt3[0]], [pt3[1]]] = np.dot(matRotation, np.array([[pt3[0]], [pt3[1]], [1]]))
    imgOut = imgRotation[int(pt1[1]):int(pt3[1]), int(pt1[0]):int(pt3[0])]


    # pt2 = list(pt2)
    # pt4 = list(pt4)
    # [[pt2[0]], [pt2[1]]] = np.dot(matRotation, np.array([[pt2[0]], [pt2[1]], [1]]))
    # [[pt4[0]], [pt4[1]]] = np.dot(matRotation, np.array([[pt4[0]], [pt4[1]], [1]]))
    # pt1 = (int(pt1[0]), int(pt1[1]))
    # pt2 = (int(pt2[0]), int(pt2[1]))
    # pt3 = (int(pt3[0]), int(pt3[1]))
    # pt4 = (int(pt4[0]), int(pt4[1]))
    # drawRect(imgRotation,pt1,pt2,pt3,pt4,(255,0,0),2)
    return imgOut


def _sen2wkt(path: str) -> str:
    """
    将 sen12 的路径名提取为多模态唯一标识符
    """
    w = os.path.basename(
        os.path.splitext(path)[0]
    ).split("_")
    
    roi = w[0] + "_" + w[1]
    region = w[3]
    patch = w[4]
    
    return str((roi, region, patch))

def _sen2modal(path: str) -> str:
    w = path.split("_")
    if w[2] == "s1":
        return "sar"
    return "opt"


# def _wkt2sen(wkt: str, modal:str) -> str:
    
#     if modal == "opt":
#         mn = "s2"
#     elif modal == "sar":
#         mn = "s1"

#     wkt = wkt.strip("()")
#     wkt = wkt.split(",")
#     for i in range(3):
#         wkt[i] = wkt[i].strip("\' ")
    
#     roi = wkt[0]
#     region = mn + "_" + wkt[1]
#     patch = roi.replace("rois", "ROIs")
#     patch += "_" + region
#     patch += "_" + wkt[2]

#     return os.path.join(patch+".png")


if not os.path.exists(img_dir2):
    os.mkdir(img_dir2)


with open(json_file, 'r') as fp:
    src_dict = json.load(fp)

cls_dict = {}
hash_map = {}
dst_dict = {}
for k in src_dict.keys():
    for fn in src_dict[k]:
        wkt = _sen2wkt(fn)
        if not wkt in cls_dict:
            cls_dict[wkt] = len(cls_dict)


for dirpath, dirnames, filenames in os.walk(img_dir1):  # 遍历文件夹
    len2 = len(dirnames)
    for k in range(len2):
        imagepath1 = str(dirpath) + '\\' + str(dirnames[k])
        imagepath2 = img_dir2 + '\\' + str(dirnames[k])
        if not os.path.exists(imagepath2):
            os.mkdir(imagepath2)
        image_list = os.listdir(imagepath1)
        k1 = 1
        for img in image_list:

            wkt = _sen2wkt(img)
            modal = _sen2modal(img)
            cur_cls = cls_dict[wkt]

            img_path1 = imagepath1 + '\\' + img
            # img1 = Image.open(img_path1).convert('L')
            img1 = cv.imread(img_path1, cv.IMREAD_GRAYSCALE)
            #  旋转
            # x1 = round(128 - 64*math.sqrt(3), 4)
            # y1 = round(128 - 64*math.sqrt(3), 4)
            # x2 = round(128 + 64*math.sqrt(3), 4)
            # y2 = round(128 + 64*math.sqrt(3), 4)
            x1 = 128 - 96
            y1 = 128 - 96
            x2 = 128 + 96
            y2 = 128 + 96
            box = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]])
            xuanzhuan = [-12, 12, -9, 9, 6, -6, 3, -3]
            for i in range(len(xuanzhuan)):
                jiaodu = xuanzhuan[i]
                imgRotation = rotateImage(img1, jiaodu, box[0], box[1], box[2], box[3])
                res = cv.resize(imgRotation, (256, 256))
                xuan_name = "c"+str(xuanzhuan[i])
                img_name = img.split('.')[0] + '_' + xuan_name + '.png'
                img_pth = imagepath2 + '\\' + img_name
                cv.imwrite(img_pth, res)

                dst_dict[img_name] = [cur_cls, modal]

            # 平移
            img2 = Image.open(img_path1).convert('L')
            # list = [0, 16, 32, 48]
            # a = random.sample(list, 3)
            for i in range(0, 64, 16):
                for j in range(0, 64, 16):
                    imgP = img2.crop((j, i, j+192, i+192))
                    imgP = np.asarray(imgP)
                    resP = cv.resize(imgP, (256, 256))
                    PingY_name = "y"+str(j) + '_' + "x"+str(i)
                    img_nameP = img.split('.')[0] + '_' + PingY_name + '.png'
                    img_pthP = imagepath2 + '\\' + img_nameP
                    cv.imwrite(img_pthP, resP)

                    dst_dict[img_nameP] = [cur_cls, modal]

with open(dst_json_path, "w") as fp:
    json.dump(dst_dict, fp, indent="\t", sort_keys=True)

